from hyde.hyde import HyDERetriever


class AdvancedHyDERetriever(HyDERetriever):
    """
    扩展的HyDE检索器，支持更多配置选项
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.cache = {}  # 缓存假设文档

    def generate_hypothetical_document(self, query: str, use_cache: bool = True) -> str:
        """
        支持缓存的假设文档生成
        """
        if use_cache and query in self.cache:
            logger.info("使用缓存的假设文档")
            return self.cache[query]

        hypothetical_doc = super().generate_hypothetical_document(query)

        if use_cache:
            self.cache[query] = hypothetical_doc

        return hypothetical_doc

    def hybrid_retrieve(self, query: str, alpha: float = 0.7, top_k: int = 3):
        """
        混合检索：结合直接查询和假设文档检索

        Args:
            query: 用户查询
            alpha: 假设文档权重（0-1）
            top_k: 检索文档数量
        """
        # 假设文档检索
        hypothetical_doc = self.generate_hypothetical_document(query)
        hyp_embedding = self.embedding_model.encode([hypothetical_doc])
        hyp_similarities = cosine_similarity(hyp_embedding, self.knowledge_embeddings)[0]

        # 直接查询检索
        query_embedding = self.embedding_model.encode([query])
        query_similarities = cosine_similarity(query_embedding, self.knowledge_embeddings)[0]

        # 混合相似度分数
        combined_similarities = alpha * hyp_similarities + (1 - alpha) * query_similarities

        # 获取最相似的文档
        top_indices = np.argsort(combined_similarities)[::-1][:top_k]

        similar_docs = []
        for idx in top_indices:
            similar_docs.append({
                'content': self.knowledge_base[idx],
                'similarity': float(combined_similarities[idx]),
                'hyde_similarity': float(hyp_similarities[idx]),
                'query_similarity': float(query_similarities[idx]),
                'index': int(idx)
            })

        return {
            'query': query,
            'hypothetical_document': hypothetical_doc,
            'retrieved_documents': similar_docs,
            'method': 'hybrid'
        }